package com.nebula.core.flink;
import com.vesoft.nebula.client.storage.StorageClient;
import com.vesoft.nebula.client.storage.data.BaseTableRow;
import com.vesoft.nebula.client.storage.data.EdgeTableRow;
import com.vesoft.nebula.client.storage.scan.ScanEdgeResult;
import com.vesoft.nebula.client.storage.scan.ScanEdgeResultIterator;
import java.util.Iterator;

import org.apache.flink.connector.nebula.statement.ExecutionOptions;

public class NebulaEdgeSource {
    ScanEdgeResultIterator iterator = null;
    Iterator<EdgeTableRow> dataIterator = null;
    StorageClient storageClient;
    ExecutionOptions executionOptions;

    public NebulaEdgeSource(StorageClient storageClient, ExecutionOptions executionOptions) {
        this.storageClient = storageClient;
        this.executionOptions = executionOptions;
    }

    public void getEdgeDataRow() {
        if (this.executionOptions.isNoColumn()) {
            this.iterator = this.storageClient.scanEdge(this.executionOptions.getGraphSpace(), this.executionOptions.getLabel(), this.executionOptions.getLimit(), this.executionOptions.getStartTime(), this.executionOptions.getEndTime(), true, true);
        } else {
            this.iterator = this.storageClient.scanEdge(this.executionOptions.getGraphSpace(), this.executionOptions.getLabel(), this.executionOptions.getFields(), this.executionOptions.getLimit(), this.executionOptions.getStartTime(), this.executionOptions.getEndTime(), true, true);
        }

    }

    public boolean hasNext() throws Exception {
        if (this.iterator == null) {
            this.getEdgeDataRow();
        }

        if (this.dataIterator == null || !this.dataIterator.hasNext()) {
            if (!this.iterator.hasNext()) {
                return false;
            }

            while(this.iterator.hasNext()) {
                ScanEdgeResult result = this.iterator.next();
                if (!result.isEmpty()) {
                    this.dataIterator = result.getEdgeTableRows().iterator();
                    break;
                }
            }
        }

        return this.dataIterator == null ? false : this.dataIterator.hasNext();
    }

    public BaseTableRow next() {
        return (BaseTableRow)this.dataIterator.next();
    }
}

